from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from utils import get_base64_img, preprocess
from langchain_core.runnables import RunnableLambda

def input_mapper_ai(screenshot_bytes: bytes, llm: ChatGoogleGenerativeAI, data_to_fill: str):
    base64_img = get_base64_img(screenshot_bytes)
    
    human_message = HumanMessage(
        content=[
            {
                "type": "text",
                "text": (
                    f"Based on the provided screenshot, map each input field in the image "
                    f"to the corresponding data field. The output should be in JSON format, "
                    f"where the key represents the index or position of the input field in the image, "
                    f"and the value is the actual data to be filled in that field. "
                    f"Given data to fill: {data_to_fill}"
                )
            },
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{base64_img}"
                }
            }
        ]
    )

    # Invoke LLM and parse the output
    mapper_chain = llm | StrOutputParser() | RunnableLambda(preprocess)
    parsed_response = mapper_chain.invoke([human_message])

    return parsed_response